// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/flat_vector.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/constant_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    template<>
    const bool *FlatVector<bool>::rawValues() const {
        POLLUX_UNSUPPORTED("rawValues() for bool is not supported");
    }

    template<>
    bool FlatVector<bool>::valueAtFast(vector_size_t idx) const {
        return bits::isBitSet(reinterpret_cast<const uint64_t *>(rawValues_), idx);
    }

    template<>
    Range<bool> FlatVector<bool>::asRange() const {
        return Range<bool>(rawValues<int64_t>(), 0, BaseVector::length_);
    }

    template<>
    void FlatVector<bool>::set(vector_size_t idx, bool value) {
        POLLUX_DCHECK_LT(idx, BaseVector::length_);
        ensureValues();
        POLLUX_DCHECK(!values_->isView());
        if (BaseVector::rawNulls_) {
            BaseVector::set_null(idx, false);
        }
        bits::setBit(reinterpret_cast<uint64_t *>(rawValues_), idx, value);
    }

    template<>
    Buffer *FlatVector<StringView>::getBufferWithSpace(
        size_t size,
        bool exactSize) {
        POLLUX_DCHECK_GE(stringBuffers_.size(), stringBufferSet_.size());

        // Check if the last buffer is uniquely referenced and has enough space.
        Buffer *buffer =
                stringBuffers_.empty() ? nullptr : stringBuffers_.back().get();
        if (buffer && buffer->unique() &&
            buffer->size() + size <= buffer->capacity()) {
            return buffer;
        }

        // Allocate a new buffer.
        const size_t newSize = exactSize ? size : std::max(kInitialStringSize, size);
        BufferPtr newBuffer = AlignedBuffer::allocate<char>(newSize, pool());
        newBuffer->setSize(0);
        addStringBuffer(newBuffer);
        return newBuffer.get();
    }

    template<>
    char *FlatVector<StringView>::getRawStringBufferWithSpace(
        size_t size,
        bool exactSize) {
        Buffer *buffer = getBufferWithSpace(size, exactSize);
        char *rawBuffer = buffer->asMutable<char>() + buffer->size();
        buffer->setSize(buffer->size() + size);
        return rawBuffer;
    }

    template<>
    void FlatVector<StringView>::prepare_for_reuse() {
        BaseVector::prepare_for_reuse();

        // Check values buffer. Keep the buffer if singly-referenced and mutable.
        // Reset otherwise.
        if (values_ && !values_->isMutable()) {
            values_ = nullptr;
            rawValues_ = nullptr;
        }

        keepAtMostOneStringBuffer();

        // Clear the StringViews to avoid referencing freed memory.
        if (rawValues_) {
            for (auto i = 0; i < BaseVector::length_; ++i) {
                rawValues_[i] = StringView();
            }
        }
    }

    template<>
    void FlatVector<StringView>::set(vector_size_t idx, StringView value) {
        POLLUX_DCHECK_LT(idx, BaseVector::length_);
        ensureValues();
        POLLUX_DCHECK(!values_->isView());
        if (BaseVector::rawNulls_) {
            BaseVector::set_null(idx, false);
        }
        if (value.isInline()) {
            rawValues_[idx] = value;
        } else {
            Buffer *buffer = getBufferWithSpace(value.size());
            char *ptr = buffer->asMutable<char>() + buffer->size();
            buffer->setSize(buffer->size() + value.size());
            memcpy(ptr, value.data(), value.size());
            rawValues_[idx] = StringView(ptr, value.size());
        }
    }

    template<>
    void FlatVector<StringView>::acquireSharedStringBuffers(
        const BaseVector *source) {
        if (!source) {
            return;
        }
        if (source->type_kind() != TypeKind::VARBINARY &&
            source->type_kind() != TypeKind::VARCHAR) {
            return;
        }
        source = source->wrapped_vector();
        switch (source->encoding()) {
            case VectorEncoding::Simple::FLAT: {
                auto *flat = source->as_unchecked<FlatVector<StringView> >();
                for (auto &buffer: flat->stringBuffers_) {
                    addStringBuffer(buffer);
                }
                break;
            }
            case VectorEncoding::Simple::CONSTANT: {
                if (!source->is_null_at(0)) {
                    auto *constant = source->as_unchecked<ConstantVector<StringView> >();
                    auto buffer = constant->getStringBuffer();
                    if (buffer != nullptr) {
                        addStringBuffer(buffer);
                    }
                }
                break;
            }

            default:
                POLLUX_UNREACHABLE(
                    "unexpected encoding inside acquireSharedStringBuffers: {}",
                    source->toString());
        }
    }

    template<>
    void FlatVector<StringView>::acquireSharedStringBuffersRecursive(
        const BaseVector *source) {
        if (!source) {
            return;
        }
        source = source->wrapped_vector();

        switch (source->encoding()) {
            case VectorEncoding::Simple::FLAT: {
                if (source->type_kind() != TypeKind::VARCHAR &&
                    source->type_kind() != TypeKind::VARBINARY) {
                    // Nothing to acquire.
                    return;
                }
                auto *flat = source->as_unchecked<FlatVector<StringView> >();
                for (auto &buffer: flat->stringBuffers_) {
                    addStringBuffer(buffer);
                }
                return;
            }

            case VectorEncoding::Simple::ARRAY: {
                acquireSharedStringBuffersRecursive(
                    source->as_unchecked<ArrayVector>()->elements().get());
                return;
            }

            case VectorEncoding::Simple::MAP: {
                acquireSharedStringBuffersRecursive(
                    source->as_unchecked<MapVector>()->mapKeys().get());
                acquireSharedStringBuffersRecursive(
                    source->as_unchecked<MapVector>()->mapValues().get());
                return;
            }

            case VectorEncoding::Simple::ROW: {
                for (auto &child: source->as_unchecked<RowVector>()->children()) {
                    acquireSharedStringBuffersRecursive(child.get());
                }
                return;
            }

            case VectorEncoding::Simple::CONSTANT: {
                // wrapped_vector can be constant vector only if the underlying type is
                // primitive.
                if (source->type_kind() != TypeKind::VARCHAR &&
                    source->type_kind() != TypeKind::VARBINARY) {
                    // Nothing to acquire.
                    return;
                }
                auto *constant_vector = source->as_unchecked<ConstantVector<StringView> >();
                if (constant_vector->is_null_at(0)) {
                    return;
                }
                auto *constant = source->as_unchecked<ConstantVector<StringView> >();
                auto buffer = constant->getStringBuffer();
                if (buffer != nullptr) {
                    addStringBuffer(buffer);
                }
                return;
            }

            case VectorEncoding::Simple::LAZY:
            case VectorEncoding::Simple::DICTIONARY:
            case VectorEncoding::Simple::SEQUENCE:
            case VectorEncoding::Simple::BIASED:
            case VectorEncoding::Simple::FUNCTION:
                POLLUX_UNREACHABLE(
                    "unexpected encoding inside acquireSharedStringBuffersRecursive: {}",
                    source->toString());
        }
    }

    template<>
    void FlatVector<StringView>::copy(
        const BaseVector *source,
        const SelectivityVector &rows,
        const vector_size_t *toSourceRow) {
        if (!rows.hasSelections()) {
            return;
        }

        // Source can be of Unknown type, in that case it should have null values.
        if (source->type_kind() == TypeKind::UNKNOWN) {
            if (source->is_constant_encoding()) {
                DKCHECK(source->is_null_at(0));
                rows.applyToSelected([&](vector_size_t row) { set_null(row, true); });
            } else {
                rows.applyToSelected([&](vector_size_t row) {
                    auto sourceRow = toSourceRow ? toSourceRow[row] : row;
                    DKCHECK(source->is_null_at(sourceRow));
                    set_null(row, true);
                });
            }
            return;
        }

        auto leaf = source->wrapped_vector()->as_unchecked<SimpleVector<StringView> >();

        if (pool_ == leaf->pool()) {
            // We copy referencing the storage of 'source'.
            copyValuesAndNulls(source, rows, toSourceRow);
            acquireSharedStringBuffers(source);
        } else {
            DecodedVector decoded(*source);
            uint64_t *rawNulls = const_cast<uint64_t *>(BaseVector::rawNulls_);
            if (decoded.may_have_nulls()) {
                rawNulls = BaseVector::mutable_raw_nulls();
            }

            size_t totalBytes = 0;
            rows.applyToSelected([&](vector_size_t row) {
                const auto sourceRow = toSourceRow ? toSourceRow[row] : row;
                if (decoded.is_null_at(sourceRow)) {
                    bits::set_null(rawNulls, row);
                } else {
                    if (rawNulls) {
                        bits::clearNull(rawNulls, row);
                    }
                    auto v = decoded.value_at<StringView>(sourceRow);
                    if (v.isInline()) {
                        rawValues_[row] = v;
                    } else {
                        totalBytes += v.size();
                    }
                }
            });

            if (totalBytes > 0) {
                auto *buffer = getRawStringBufferWithSpace(totalBytes);
                rows.applyToSelected([&](vector_size_t row) {
                    const auto sourceRow = toSourceRow ? toSourceRow[row] : row;
                    if (!decoded.is_null_at(sourceRow)) {
                        auto v = decoded.value_at<StringView>(sourceRow);
                        if (!v.isInline()) {
                            memcpy(buffer, v.data(), v.size());
                            rawValues_[row] = StringView(buffer, v.size());
                            buffer += v.size();
                        }
                    }
                });
            }
        }

        if (auto stringVector = source->as<SimpleVector<StringView> >()) {
            if (auto ascii = stringVector->isAscii(rows, toSourceRow)) {
                setIsAscii(ascii.value(), rows);
            } else {
                // ASCII-ness for the 'rows' is not known.
                ensureIsAsciiCapacity();
                // If we arent All ascii, then invalidate
                // because the remaining selected rows might be ascii
                if (!asciiInfo.isAllAscii()) {
                    invalidateIsAscii();
                } else {
                    asciiInfo.writeLockedAsciiComputedRows()->deselect(rows);
                }
            }
        }
    }

    // For strings if backing memory pool is not the same as the vector pool, we
    // need to perform a deep copy and reconstruct the string views against the
    // updated stringBuffers.
    template<>
    VectorPtr FlatVector<StringView>::copy_preserve_encodings(
        pollux::memory::MemoryPool *pool) const {
        const auto allocPool = pool ? pool : BaseVector::pool_;
        // If the backing memory pool is the same as the vector pool
        // we can do a shallow copy as string buffers can be shared.
        if (pool == BaseVector::pool_) {
            return std::make_shared<FlatVector<StringView> >(
                allocPool,
                BaseVector::type_,
                AlignedBuffer::copy(allocPool, BaseVector::nulls_),
                BaseVector::length_,
                AlignedBuffer::copy(allocPool, values_),
                std::vector<BufferPtr>(stringBuffers_),
                SimpleVector<StringView>::stats_,
                BaseVector::distinctValueCount_,
                BaseVector::nullCount_,
                SimpleVector<StringView>::isSorted_,
                BaseVector::representedByteCount_,
                BaseVector::storageByteCount_);
        } else {
            size_t totalBytes = 0;
            auto newValuesBuffer =
                    AlignedBuffer::allocate<StringView>(BaseVector::size(), allocPool);
            auto *rawCopyValues = newValuesBuffer->asMutable<StringView>();
            // Copy non Null StringViews to value buffer.
            for (vector_size_t i = 0; i < BaseVector::size(); i++) {
                if (!BaseVector::is_null_at(i)) {
                    auto v = value_at(i);
                    if (v.isInline()) {
                        rawCopyValues[i] = v;
                    } else {
                        totalBytes += v.size();
                    }
                }
            }

            std::vector<BufferPtr> newStringBuffers;
            if (totalBytes > 0) {
                newStringBuffers.emplace_back(
                    AlignedBuffer::allocate<char>(totalBytes, allocPool));
                char *rawBuffer = newStringBuffers.back()->asMutable<char>();

                for (vector_size_t i = 0; i < BaseVector::size(); i++) {
                    if (!BaseVector::is_null_at(i)) {
                        auto v = value_at(i);
                        if (!v.isInline()) {
                            memcpy(rawBuffer, v.data(), v.size());
                            rawCopyValues[i] = StringView(rawBuffer, v.size());
                            rawBuffer += v.size();
                        }
                    }
                }
            }

            return std::make_shared<FlatVector<StringView> >(
                allocPool,
                BaseVector::type_,
                AlignedBuffer::copy(allocPool, BaseVector::nulls_),
                BaseVector::length_,
                newValuesBuffer,
                std::move(newStringBuffers),
                SimpleVector<StringView>::stats_,
                BaseVector::distinctValueCount_,
                BaseVector::nullCount_,
                SimpleVector<StringView>::isSorted_,
                BaseVector::representedByteCount_,
                BaseVector::storageByteCount_);
        }
    }

    // For strings, we also verify if they point to valid memory locations inside
    // the string buffers.
    template<>
    void FlatVector<StringView>::validate(
        const VectorValidateOptions &options) const {
        SimpleVector<StringView>::validate(options);
        auto byteSize = BaseVector::byteSize<StringView>(BaseVector::size());
        if (byteSize == 0) {
            return;
        }
        POLLUX_CHECK_NOT_NULL(values_);
        POLLUX_CHECK_GE(values_->size(), byteSize);
        auto rawValues = values_->as<StringView>();

        for (auto i = 0; i < BaseVector::length_; ++i) {
            if (is_null_at(i)) {
                continue;
            }
            auto stringView = rawValues[i];
            if (!stringView.isInline()) {
                bool isValid = false;
                for (const auto &buffer: stringBuffers_) {
                    auto start = buffer->as<char>();
                    if (stringView.data() >= start &&
                        stringView.data() < start + buffer->size()) {
                        isValid = true;
                        break;
                    }
                }
                POLLUX_CHECK(
                    isValid,
                    "String view at idx {} points outside of the string buffers",
                    i);
            }
        }
    }
} // namespace kumo::pollux
